"""
示例 4: 创建磁场阵列

演示如何：
1. 创建磁场阵列
2. 设置磁场分量
3. 创建不同类型的磁场分布
"""

import sys
import os
from math import sin, cos, pi, exp

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from SIMION.PA import PA


def create_uniform_magnetic_field():
    """创建均匀磁场"""
    
    print("\n" + "="*60)
    print("创建均匀磁场")
    print("="*60 + "\n")
    
    # 创建磁场阵列
    # 注意: field_type='magnetic'
    pa = PA(
        nx=50,
        ny=50,
        nz=50,
        symmetry='planar',
        field_type='magnetic',  # 磁场类型
        ng=100,                 # ng 缩放因子
        dx_mm=1.0,
        dy_mm=1.0,
        dz_mm=1.0
    )
    
    print(f"创建磁场阵列: {pa.nx()} x {pa.ny()} x {pa.nz()}")
    print(f"场类型: {pa.field_type()}")
    print(f"ng 因子: {pa.ng()}\n")
    
    # 设置均匀磁场（沿 Z 方向）
    # 磁场强度: 0.1 Tesla = 1000 Gauss
    Bz = 1000.0  # Gauss
    
    print(f"设置均匀磁场: Bz = {Bz} Gauss")
    print("填充磁场数据...")
    
    for z in range(pa.nz()):
        for y in range(pa.ny()):
            for x in range(pa.nx()):
                # field(x, y, z, Bx, By, Bz)
                pa.field(x, y, z, 0.0, 0.0, Bz)
    
    # 验证
    Bx, By, Bz_read = pa.field(25, 25, 25)
    print(f"\n验证中心点磁场:")
    print(f"  坐标: (25, 25, 25)")
    print(f"  Bx = {Bx} Gauss")
    print(f"  By = {By} Gauss")
    print(f"  Bz = {Bz_read} Gauss")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "uniform_magnetic.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60 + "\n")
    
    return pa


def create_gradient_magnetic_field():
    """创建梯度磁场"""
    
    print("\n" + "="*60)
    print("创建梯度磁场")
    print("="*60 + "\n")
    
    pa = PA(
        nx=60,
        ny=60,
        nz=60,
        symmetry='planar',
        field_type='magnetic',
        ng=100,
        dx_mm=1.0,
        dy_mm=1.0,
        dz_mm=1.0
    )
    
    print(f"创建磁场阵列: {pa.nx()} x {pa.ny()} x {pa.nz()}\n")
    
    # 创建沿 Z 方向线性增加的磁场
    B_min = 500.0   # 最小磁场强度 (Gauss)
    B_max = 1500.0  # 最大磁场强度 (Gauss)
    
    print(f"磁场梯度:")
    print(f"  Z=0 处: {B_min} Gauss")
    print(f"  Z={pa.nz()-1} 处: {B_max} Gauss")
    print(f"  梯度: {(B_max-B_min)/(pa.nz()-1):.2f} Gauss/网格\n")
    
    print("填充磁场数据...")
    
    for z in range(pa.nz()):
        # 磁场强度随 Z 线性变化
        Bz = B_min + (B_max - B_min) * z / (pa.nz() - 1)
        
        for y in range(pa.ny()):
            for x in range(pa.nx()):
                pa.field(x, y, z, 0.0, 0.0, Bz)
    
    # 验证不同位置
    print(f"验证磁场分布:")
    for z_test in [0, pa.nz()//2, pa.nz()-1]:
        Bx, By, Bz_read = pa.field(30, 30, z_test)
        print(f"  Z={z_test}: Bz = {Bz_read:.2f} Gauss")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "gradient_magnetic.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60 + "\n")
    
    return pa


def create_cylindrical_magnetic_field():
    """创建圆柱对称磁场（螺线管）"""
    
    print("\n" + "="*60)
    print("创建圆柱对称磁场（螺线管）")
    print("="*60 + "\n")
    
    pa = PA(
        nx=40,                     # 径向
        ny=80,                     # 轴向
        nz=1,
        symmetry='cylindrical',
        field_type='magnetic',
        ng=100,
        dx_mm=1.0,
        dy_mm=1.0
    )
    
    print(f"创建磁场阵列: {pa.nx()} x {pa.ny()} (圆柱对称)\n")
    
    # 模拟螺线管内部的磁场
    # 内部: 均匀磁场
    # 外部: 磁场衰减
    
    solenoid_radius = 20  # 螺线管半径
    B_inside = 1000.0     # 内部磁场强度 (Gauss)
    
    print(f"螺线管参数:")
    print(f"  半径: {solenoid_radius} 网格 ({solenoid_radius*pa.dx_mm()} mm)")
    print(f"  内部磁场: {B_inside} Gauss\n")
    
    print("填充磁场数据...")
    
    for y in range(pa.ny()):
        for r in range(pa.nx()):
            if r <= solenoid_radius:
                # 螺线管内部: 均匀轴向磁场
                Br = 0.0
                Bz = B_inside
            else:
                # 螺线管外部: 磁场衰减
                # 简化模型: 1/r^2 衰减
                decay_factor = (solenoid_radius / r) ** 2
                Br = 0.0
                Bz = B_inside * decay_factor * 0.1
            
            # 注意: 在圆柱对称中，Y 是轴向，X 是径向
            # field(r, z, 0, Br, Bz, Bphi)
            pa.field(r, y, 0, Br, Bz, 0.0)
    
    # 验证不同半径处的磁场
    print(f"验证磁场分布 (Y={pa.ny()//2}):")
    for r_test in [10, 20, 30]:
        Br, Bz_read, Bphi = pa.field(r_test, pa.ny()//2, 0)
        print(f"  R={r_test}: Br={Br:.2f}, Bz={Bz_read:.2f} Gauss")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "solenoid_magnetic.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60 + "\n")
    
    return pa


def create_quadrupole_magnetic_field():
    """创建四极磁场"""
    
    print("\n" + "="*60)
    print("创建四极磁场")
    print("="*60 + "\n")
    
    pa = PA(
        nx=60,
        ny=60,
        nz=40,
        symmetry='planar',
        field_type='magnetic',
        ng=100,
        dx_mm=0.5,
        dy_mm=0.5,
        dz_mm=0.5
    )
    
    print(f"创建磁场阵列: {pa.nx()} x {pa.ny()} x {pa.nz()}\n")
    
    # 四极磁场: Bx = k*y, By = k*x
    # 其中 k 是梯度常数
    
    k = 20.0  # 梯度常数 (Gauss/网格)
    center_x = pa.nx() // 2
    center_y = pa.ny() // 2
    
    print(f"四极磁场参数:")
    print(f"  中心: ({center_x}, {center_y})")
    print(f"  梯度常数: {k} Gauss/网格\n")
    
    print("填充磁场数据...")
    
    for z in range(pa.nz()):
        for y in range(pa.ny()):
            for x in range(pa.nx()):
                # 相对于中心的坐标
                dx = x - center_x
                dy = y - center_y
                
                # 四极磁场公式
                Bx = k * dy
                By = k * dx
                Bz = 0.0
                
                pa.field(x, y, z, Bx, By, Bz)
    
    # 验证四个象限
    print(f"验证四极磁场 (Z={pa.nz()//2}):")
    offset = 10
    test_points = [
        (center_x + offset, center_y, "右侧"),
        (center_x - offset, center_y, "左侧"),
        (center_x, center_y + offset, "上侧"),
        (center_x, center_y - offset, "下侧")
    ]
    
    for x, y, label in test_points:
        Bx, By, Bz = pa.field(x, y, pa.nz()//2)
        print(f"  {label} ({x},{y}): Bx={Bx:.2f}, By={By:.2f} Gauss")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "quadrupole_magnetic.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60 + "\n")
    
    return pa


if __name__ == "__main__":
    print("\n" + "="*60)
    print("SIMION 磁场阵列创建示例")
    print("="*60)
    
    # 1. 均匀磁场
    pa1 = create_uniform_magnetic_field()
    
    # 2. 梯度磁场
    pa2 = create_gradient_magnetic_field()
    
    # 3. 螺线管磁场
    pa3 = create_cylindrical_magnetic_field()
    
    # 4. 四极磁场
    pa4 = create_quadrupole_magnetic_field()
    
    print("\n" + "="*60)
    print("所有磁场阵列创建完成!")
    print("="*60)
    print("\n生成的文件:")
    print("  - uniform_magnetic.pa# : 均匀磁场")
    print("  - gradient_magnetic.pa# : 梯度磁场")
    print("  - solenoid_magnetic.pa# : 螺线管磁场")
    print("  - quadrupole_magnetic.pa# : 四极磁场")
    print("\n可以在 SIMION 中打开这些文件进行粒子模拟\n")

